package com.yuxl.ai.embedding.bailian.service;

import com.google.common.collect.Lists;
import com.yuxl.ai.embedding.bailian.domain.BailianEmbeddingData;
import com.yuxl.ai.embedding.bailian.domain.BailianEmbeddingReq;
import com.yuxl.ai.embedding.bailian.domain.BailianEmbeddingRes;
import com.yuxl.ai.embedding.config.EmbeddingProperties;
import com.yuxl.ai.embedding.service.EmbeddingApi;
import com.yuxl.common.core.utils.HttpUtil;
import com.yuxl.common.core.utils.ObjectUtil;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.http.client.reactive.ReactorClientHttpConnector;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.netty.http.client.HttpClient;

import java.util.List;

/**
 * @Author : yuxl
 * @create 2025/7/10 下午5:08
 */
@Slf4j
@Service
@ConditionalOnProperty(prefix = "ai.embedding", name = "type", havingValue = "bailian")
public class BailianEmbeddingApi implements EmbeddingApi {
    @Resource
    private EmbeddingProperties embeddingProperties;
    @Resource
    private ObjectUtil objectUtil;
    public Mono<float[]> embedding(String text){
        return embedding(List.of(text)).map(r->r.get(0));
    }
    public Mono<List<float[]>> embedding(List<String> texts){
        int batchSize = embeddingProperties.getMaxRows();
        List<Mono<List<float[]>>> tasks = Lists.partition(texts, batchSize)
                .stream()
                .map(textList->{
                    BailianEmbeddingReq body = new BailianEmbeddingReq()
                            .setModel(embeddingProperties.getModel())
                            .setInput(textList)
                            ;
                    return api(body).map(res-> res.getData().stream().map(BailianEmbeddingData::getEmbedding).toList());
                }).toList();
        return Flux.concat(tasks)
                .flatMap(Flux::fromIterable)
                .collectList();
    }

    public Mono<BailianEmbeddingRes> api(BailianEmbeddingReq body) {
        HttpClient httpClient = HttpUtil.httpClientSetTimeout(1000, Long.MAX_VALUE, Long.MAX_VALUE);
        String b = objectUtil.writeValueAsString(body);
        log.debug("embedding-api=>body:{}",b);
        return WebClient
                .builder()
                .codecs(configurer ->
                        configurer.defaultCodecs().maxInMemorySize(20 * 1024 * 1024) // 设置为 20MB
                )
                .clientConnector(new ReactorClientHttpConnector(httpClient))
                .baseUrl(embeddingProperties.getBaseUrl())
                .build()
                .post()
                .header("content-type", "application/json")
                .header("Authorization", "Bearer " + embeddingProperties.getApiKey())
                .bodyValue(body)
                .retrieve()
                .bodyToMono(String.class)
                .map(r -> {
                    log.debug("embedding-api=>res:{}",r);
                    return objectUtil.readValue(r, BailianEmbeddingRes.class);
                })
                ;
    }
}
